'''
data_root=/mnt/nas/datasets/diction
save_root=/mnt/nas/shengjie/puthere_data
date=0821
categories=("coat" "leather" "sweater")

date=0916
categories=("trench" "paike" "leafur")

date=0918
categories=("dress" "slipdress" "feather")

clo_num=100
get_clo_dir=os.path.join( data_dir, $category$date_img_clo )
names get from get_clo_dir/names.txt, sift ends with ".jpg", random sift limit $clo_num
unifiy filename from $category_clo_img_1 ~ $category_clo_img_100 .jpg
img save to save_root/$category_clo

human_num=100
get_human_dir=os.path.join( data_dir, $category$date_img_human )
names get from get_human_dir/names.txt, sift ends with ".jpg", random sift limit $human_num
unifiy filename from $category_human_img_1 ~ $category_human_img_100 .jpg
img save to save_root/$category_human
'''
import os
import random
from shutil import copyfile
from PIL import Image

def get_jpg_names(names_path, limit):
    with open(names_path, "r") as f:
        all_names = [line.strip() for line in f if line.strip()]
    jpg_names = [name for name in all_names if name.lower().endswith('.jpg')]
    random.shuffle(jpg_names)
    return jpg_names[:limit]

def save_images(src_dir, names, dst_dir, prefix):
    os.makedirs(dst_dir, exist_ok=True)
    for idx, name in enumerate(names, 1):
        src_path = os.path.join(src_dir, name)
        ext = os.path.splitext(name)[1]
        dst_name = f"{prefix}_{idx}{ext}"
        dst_path = os.path.join(dst_dir, dst_name)
        # 统一保存为jpg格式
        try:
            img = Image.open(src_path)
            img = img.convert("RGB")
            img.save(dst_path, "JPEG")
        except Exception as e:
            print(f"Error processing {src_path}: {e}")

def process_category(category, date, data_root, save_root, clo_num=100, human_num=100):
    # 服装图片
    clo_dir = os.path.join(data_root, f"{category}{date}_img_clo")
    clo_names_path = os.path.join(clo_dir, "names.txt")
    clo_names = get_jpg_names(clo_names_path, clo_num)
    save_images(clo_dir, clo_names, os.path.join(save_root, f"{category}_clo"), f"{category}_clo_img")

    # 人体图片
    human_dir = os.path.join(data_root, f"{category}{date}_img_human")
    human_names_path = os.path.join(human_dir, "names.txt")
    human_names = get_jpg_names(human_names_path, human_num)
    save_images(human_dir, human_names, os.path.join(save_root, f"{category}_human"), f"{category}_human_img")

if __name__ == "__main__":
    data_root = "/mnt/nas/datasets/diction"
    save_root = "/mnt/nas/shengjie/puthere_data"
    configs = [
        # (date, [categories])
        ("0821", ["coat", "leather", "sweater"]),
        ("0916", ["trench", "paike", "leafur"]),
        ("0918", ["dress", "slipdress", "feather"]),
    ]
    clo_num = 100
    human_num = 100

    # INSERT_YOUR_CODE
    import os

    # 检查所有的文件夹和names.txt文件是否存在
    missing_items = []
    for date, categories in configs:
        for category in categories:
            # 检查服装图片文件夹和names.txt
            clo_dir = os.path.join(data_root, f"{category}{date}_img_clo")
            clo_names_path = os.path.join(clo_dir, "names.txt")
            if not os.path.isdir(clo_dir):
                missing_items.append(f"Missing directory: {clo_dir}")
            if not os.path.isfile(clo_names_path):
                missing_items.append(f"Missing file: {clo_names_path}")

            # 检查人体图片文件夹和names.txt
            human_dir = os.path.join(data_root, f"{category}{date}_img_human")
            human_names_path = os.path.join(human_dir, "names.txt")
            if not os.path.isdir(human_dir):
                missing_items.append(f"Missing directory: {human_dir}")
            if not os.path.isfile(human_names_path):
                missing_items.append(f"Missing file: {human_names_path}")

    if missing_items:
        print("❌ The following required files or directories are missing:")
        for item in missing_items:
            print("   ", item)
        print("Please check your data_root and ensure all files exist before running the script.")
        exit(1)
    else:
        print("✅ All required directories and names.txt files exist. Proceeding...")

    for date, categories in configs:
        for category in categories:
            print(f"Processing {category} ({date}) ...")
            process_category(category, date, data_root, save_root, clo_num, human_num)
    print("✅ All categories processed and images saved.")


